import copy

import torch
import torch.nn as nn
import transformer as tf
from torch.autograd import Variable

c = copy.deepcopy

def mask_model(source_vocab, target_vocab, N=8,
               d_model=512, d_ff=1024, head=8, dropout=0.1):
    # source_vocab是源数据特征（词汇）总数，target_vocab是目标数据特征（词汇）总数，
    # N是编码器和解码器堆叠数，d_model是词向量映射维度，d_ff前鐀全连接网络中变换矩阵的维度
    # head是多头注意力结构中的多头数，dropout是置零比率
    # 首先得到一个深度拷贝命令，接下来很多结构都需要进行深度拷贝
    # 来保证他们彼此之间相互独立，不受干扰
    c = copy.deepcopy

    # 实例化了多头注意力类，得到对象attn
    attn = tf.MultiHeadedAttention(head, d_model)

    # 然后实例化前鐀全连接类，得到对象ff
    ff = tf.PositionwiseFeedForward(d_model, d_ff, dropout)

    # 实例化位置编码类，得到对象position
    position = tf.PositionalEncoding(d_model, dropout)

    # 根据结构图，最外层是EncoderDecoder,在EncoderDecoder中，
    # 分别是是编码器层，解码器层，源数据Embedding层和位置编码组成的有序结构
    # 目标数据Embedding层和位置编码组成的有序结构，以及类别生成器层，
    # 在编码器层中有attention子层以及前鐀全连接子层
    # 在解码器层中有两个attention子层以及前鐀全连接层
    model = tf.EncoderDecoder(
        tf.Encoder(tf.EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        tf.Decoder(tf.DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        nn.Sequential(tf.Embeddings(d_model, source_vocab), c(position)),
        nn.Sequential(tf.Embeddings(d_model, target_vocab), c(position)),
        tf.Generator(d_model, target_vocab)
    )
    # 模型结构完成后，接下来就是初始化模型中的参数，比如线性层中的变换矩阵
    # 这里一但判断参数的维度大于1，则会将其初始化成一个服从均匀分布的矩阵，
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform(p)
    return model

source_vocab = 1000
target_vocab = 1000
d_model = 512
N = 6
head = 8
dropout = 0.2

model = mask_model(source_vocab, target_vocab, N, d_model=d_model, dropout=dropout, head=head)
print(model)

source = target = Variable(torch.LongTensor(
    [
        [100, 2, 421, 508],
        [491, 998, 1, 221]
    ]
))

source_mask = target_mask = Variable(torch.zeros(head, 4, 4))

ed_result = model(source, target, source_mask, target_mask)
print(ed_result)
print(ed_result.shape)
